
import numpy as np
import sophon.sail as sail
from easydict import EasyDict as edict
import cv2
import time
import datetime
import logging
from pathlib import Path
from PIL import Image
import yaml
import argparse
import shutil

from utils import set_logger
from camera import CameraManager
from communicate import Communicate, Message, ImageSender

import time

from detector import Detector
from watchdog import WatchDog

gvars = edict({
    'clean_interval': None,
    'save_image': True,
    'conf_thres': None,
    'auto_save': None,

    'model': None,
    'count_toggle': None
})

def process_msg(msg):
    type, instr, data = msg
    if type == 'error':
        logging.error("receive error data: %s, %s" % (instr, data))
    elif type == 'response':
        logging.info('receive response data: %s, %s' % (instr, data))
    elif type == 'control':
        if instr == 'configure':
            logging.info('receive control data: %s, %s' % (instr, data))
            conf = dict(tuple(s.split('=') for s in data.split('&')))
            gvars.conf_thres = float(conf['t'])
            gvars.model.set_conf_thres(gvars.conf_thres)
            gvars.clean_interval = float(conf['c'])
            gvars.auto_save = bool(int(conf['a']))
            gvars.save_image = bool(int(conf['s']))
        if instr == 'toggle':
            logging.info('receive control data: %s, %s' % (instr, data))
            conf = dict(tuple(s.split('=') for s in data.split('&')))
            gvars.count_toggle[conf['c']] = bool(int(conf['r']))

def main(args):
    logging.info('main start')
    image_save_dir = Path(args.image_save_dir)
    image_save_dir.mkdir(parents=True, exist_ok=True)
    gvars.clean_interval = args.clean_interval
    gvars.conf_thres = args.conf_thres


    detector = Detector(
        args.bmodel, args.conf_thres, 0.4,
        min_box=args.get('min_box', (0, 0)),
        extra_scales=args.get('extra_scales', ())
    )
    logging.info('detectot build finished.')
    gvars.detector = detector

    cams = CameraManager(args.input)
    gvars.count_toggle = {k:True for k in args.input}

    if args.enable_com:
        com = Communicate(args.server_ip, args.server_port)
        com.handshake(args.box_id, args.box_code)

        image_sender = ImageSender(args.image_server_url)

    watch_dog = WatchDog(60, name='main_dog')
    last_clean_time = time.time()
    last_det_time = None
    time.sleep(1)
    while True:
        # process received msgs
        while args.enable_com and not com.msgs.empty():
            msg = com.msgs.get()
            process_msg(msg)
        
        frames = cams.read_all_cameras()
        watch_dog.feed()
        t = datetime.datetime.now()
        t = t.replace(second=t.second // 10 * 10)
        if last_det_time is not None and (t - last_det_time).total_seconds() <= 10:
            time.sleep(0.5)
            continue
        last_det_time = t
        timestamp = t.strftime('%Y%m%d%H%M%S')
        
        for cam_id, frame in frames.items():
            if not gvars.count_toggle[cam_id]:
                continue
            margins = args.margins.get(cam_id, None)

            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            logging.info('detecting: %s' % cam_id)
            _t = time.time()
            boxes, confs = detector(rgb_frame, margins=margins)
            logging.info('detected: %d, consuming: %.3f secs' % (len(boxes), time.time() - _t))

            # save image to local disk
            if margins is not None:
                cv2.rectangle(
                    rgb_frame,
                    (margins[0], margins[1]),
                    (3840-margins[2], 2160-margins[3]),
                    (0, 255, 0)
                )
            for box in boxes:
                box = box.astype(np.int)
                cv2.rectangle(rgb_frame, (box[0], box[1]), (box[2], box[3]), (255, 0, 0))
            out_file = image_save_dir / (cam_id + '_' + timestamp + '.jpg')
            logging.info('save image: %s' % out_file)
            Image.fromarray(rgb_frame).save(str(out_file))

            if args.enable_com:
                # send image to server
                image_sender.send(args.box_id, args.box_id, out_file)
                # send person counts
                com.send_count(cam_id, timestamp, len(boxes))

        time.sleep(args.sample_interval)
        _toc = time.time()
        if _toc - last_clean_time > gvars.clean_interval:
            logging.info('time to remove histories')
            shutil.rmtree(str(image_save_dir))
            image_save_dir.mkdir(parents=True, exist_ok=True)            
            last_clean_time = _toc
    
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('config_file', default='./main.yaml', type=str, nargs='?', help='.yaml config file')
    cmd_args = parser.parse_args()

    args = yaml.load(open(cmd_args.config_file, 'r'), Loader=yaml.FullLoader)
    args = edict(args)

    set_logger()
    main(args)
